import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import json

# 读取配置文件
with open('../Input.json', 'r') as f:
    config = json.load(f)

# 从配置文件中读取网格分辨率参数
parm = config["MultiGrid_Params"]["grid_resolution"]

# 创建一个大图，包含多个子图
fig = plt.figure(figsize=(15, 10))

if config["Dimension"] == 2:
    # 遍历每个网格分辨率
    for idx, resolution in enumerate(parm):
        # 读取数据
        filename = f'test_parm{resolution}.txt'
        data = np.loadtxt(filename)

        # 提取x, y, z坐标
        x = data[:, 0]
        y = data[:, 1]
        z = data[:, 2]

        # 创建网格
        # 由于部分点被移除，需要重新构造网格
        m = resolution - 1# 网格分辨率（包括边界点）
        x_grid = np.linspace(0, 1, m)
        y_grid = np.linspace(0, 1, m)
        x, y = np.meshgrid(x_grid, y_grid)

        # 将z值填充到网格中
        z_grid = np.full((m, m), np.nan)  # 初始化网格为NaN
        for i in range(len(data)):
            xi = int(np.round(data[i, 0] * (m - 1)))  # 计算x索引
            yi = int(np.round(data[i, 1] * (m - 1)))  # 计算y索引
            z_grid[yi, xi] = data[i, 2]  # 填充z值

        # 添加子图
        ax = fig.add_subplot(2, 2, idx + 1, projection='3d')  # 2行2列的子图布局

        # 绘制表面图
        surf = ax.plot_surface(x, y, z_grid, cmap='viridis')

        fig.colorbar(surf, ax=ax, shrink=0.6, aspect=10)
        # 设置标签和标题
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('U')
        ax.set_title(f'Grid Resolution: {resolution}x{resolution}')

    # 调整布局
    plt.tight_layout()
    plt.savefig(f'../figure/test_2d_{config["source_term"]}_{config["MultiGrid_Params"]["type"]}_{config["BoundaryConditions"]["type"]}_{config["MultiGrid_Params"]["relative_accuracy"]}_{config["MultiGrid_Params"]["restriction"]}_{config["MultiGrid_Params"]["interpolation"]}.png')

if config["Dimension"] == 1:
    # 遍历每个网格分辨率
    for idx, resolution in enumerate(parm):
        # 读取数据
        filename = f'test_parm{resolution}.txt'
        data = np.loadtxt(filename)

        # 提取x, y坐标
        x = data[:, 0]
        y = data[:, 1]

        # 添加子图
        ax = fig.add_subplot(2, 2, idx + 1)  # 2行2列的子图布局

        # 绘制折线图
        ax.plot(x, y, label=f'Resolution: {resolution}', color='b', marker='o')

        # 设置标签和标题
        ax.set_xlabel('X')
        ax.set_ylabel('U')
        ax.set_title(f'Grid Resolution: {resolution}')
        ax.legend()

    # 调整布局
    plt.tight_layout()
    plt.savefig(f'../figure/test_1d_{config["source_term"]}_{config["MultiGrid_Params"]["type"]}_{config["BoundaryConditions"]["type"]}_{config["MultiGrid_Params"]["relative_accuracy"]}_{config["MultiGrid_Params"]["restriction"]}_{config["MultiGrid_Params"]["interpolation"]}.png')

# 显示图形
plt.show()